package com.shuai.robot.controller;

import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.LoadingCache;
import com.shuai.common.constant.AuthenticationConstant;
import com.shuai.common.constant.ErrorCode;
import com.shuai.common.exception.BusinessException;
import com.shuai.common.rps.BaseResponse;
import com.shuai.common.utils.MyJWTUtil;
import com.shuai.common.utils.ResultUtils;
import com.shuai.common.utils.SnowflakeIdGenerator;
import com.shuai.robot.dto.StreamChatRequest;
import com.shuai.robot.service.StreamChatService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.validation.Valid;
import java.io.IOException;
import java.util.concurrent.TimeUnit;

@RestController
@RequestMapping("/sse")
@Slf4j
public class SseStreamChatController {


    @Autowired
    private StreamChatService streamChatService;


    /**
     * SseEmitter缓存
     */
    private LoadingCache<Long, SseEmitter> sseEmitterMap = Caffeine.newBuilder().initialCapacity(1024)
            // 手动设置2分钟缓存过期，一次流式请求不可能超过2分钟
            .expireAfterAccess(2L, TimeUnit.MINUTES)
            //缓存填充策略
            .build(sseId -> new SseEmitter(0L));

    /**
     * 获取SSE连接，返回sseID给前端
     *
     * @return
     */
    @GetMapping("/chat/getSseEmitter")
    public SseEmitter getSseEmitter() {

        // 默认30秒超时,设置为0L则永不超时
        SseEmitter sseEmitter = new SseEmitter(0L);

        // 生成sseID 并且通过SseEmitter传递给前端，后续前端通过sseID来发送消息（实现双向通讯）
        long sseEmitterId = SnowflakeIdGenerator.nextId();

        try {
            sseEmitter.send(SseEmitter.event().id("SET_CONNECT").data(sseEmitterId).reconnectTime(3000));

        } catch (IOException e) {
            log.error("获取SSE连接失败！");
            return null;
        }
        // 将SseEmitter添加到连接缓存中
        sseEmitterMap.put(sseEmitterId, sseEmitter);
        log.info("获取SSE连接成功！{}",sseEmitterId);
        return sseEmitter;
    }


    /**
     * 多轮流式对话
     *
     * @param req
     * @return
     */
    @PostMapping(path = "/chat/streamSessionChat")
    public BaseResponse<Boolean> streamSessionChat(@RequestBody @Valid StreamChatRequest req, @RequestHeader(AuthenticationConstant.AUTHENTICATION_HEADER) String authentication) {

        // TODO: 如果用户上传了apiKey和apiSecret，则使用用户的，否则使用本系统默认配置的
        // 拿到用户ID
        String userId = (String) MyJWTUtil.getPayLoadItem(authentication, AuthenticationConstant.PAYLOAD_USERID);

        // 获取指定的sseEmitter, 将响应信息通过sseEmitter发送出去
        SseEmitter sseEmitter = sseEmitterMap.get(req.getSseEmitterId());
        if (sseEmitter == null) {
            throw new BusinessException(ErrorCode.OPERATION_ERROR, "sseEmitter连接失效！");
        }

        streamChatService.streamChat(req.getSessionId(), userId, req.getQuestion(), sseEmitter);

        // 清除缓存
        sseEmitterMap.invalidate(req.getSseEmitterId());
        return ResultUtils.success(Boolean.TRUE);
    }
}
